package jtg.generator;

import com.microsoft.z3.*;
import jtg.graphics.SootCFG;
import jtg.parser.Z3Parser;
import jtg.state.PathState;
import jtg.trie.TrieNode;
import jtg.trie.TrieTree;
import jtg.visualizer.Visualizer;
import soot.Body;
import soot.Local;
import soot.Unit;
import soot.Value;
import soot.jimple.internal.*;
import soot.toolkits.graph.UnitGraph;

import java.io.File;
import java.util.*;

public class StateCoverageGenerator extends Generator{

    Set<Unit> visitedNode = new HashSet<>();

    public StateCoverageGenerator(String className, String methodName) {
        super(className, methodName);
        String defaultClsPath = System.getProperty("user.dir") + File.separator + "target" + File.separator + "classes";
        new PrimeCoverageGenerator(defaultClsPath, className, methodName);
    }

    public StateCoverageGenerator(String classPath, String className, String methodName) {
        super(classPath, className, methodName);
    }

    @Override
    void drawCFG() {
        drawCFG(clsName + "." + mtdName, false);
    }

    public void drawCFG(String graphName, boolean indexLabel) {
        Visualizer.printCFGDot(graphName, ug, indexLabel);
    }

    private List<Local> getJVars() {
        //Jimple自身增加的Locals，不是被测代码真正的变量
        ArrayList<Local> jimpleVars = new ArrayList<Local>();
        for (Local l : body.getLocals()) {
            if (l.toString().startsWith("$")) jimpleVars.add(l);
        }
        return jimpleVars;
    }

    private List<Local> getParameters() {
        return new ArrayList<>(body.getParameterLocals());
    }

    public List<String> generate() {

        ArrayList<String> testSet = null;

        System.out.println("============================================================================");
        System.out.println("Generating test case inputs for method: " + clsName + "." + mtdName + "()");
        System.out.println("============================================================================");
        try {
            testSet = new ArrayList<String>();
            Context context = new Context();
            Solver solver = context.mkSolver();

            if(ug.getHeads().size()>1) throw new Exception("The number of heads is more than 1");
            Unit h = ug.getHeads().get(0);
            // path = ug.getExtendedBasicBlockPathBetween(h, t);
            List<PathState> pathStates = initPathState(context, h);
            calculateWithBFS(context, ug, pathStates, testSet);
            int allNodes = ug.size();
            int coveredNodes = visitedNode.size();
            System.out.println("The number of all nodes is: " + allNodes + ", the number of covered nodes is: " + coveredNodes);


        } catch (Exception e) {
            System.err.println("Error in generating test cases: ");
            System.err.println(e.toString());
            e.printStackTrace();
        }
        if (!testSet.isEmpty()) {
            System.out.println("");
            System.out.println("The generated test case inputs:");
            int count = 1;
            for (String tc : testSet) {
                System.out.println("( " + count + " ) " + tc.toString());
                count++;
            }
        }
        return testSet;
    }

    public List<PathState> initPathState(Context context, Unit h){
        List<PathState> pathStates = new ArrayList<>();
        // 创建初始状态
        PathState initialState = new PathState(h);
        for(Local local : getParameters()){
            String parName = local.getName();
            initialState.params.put(parName, context.mkIntConst(parName));
            initialState.paramLatestMap.put(parName,0);
        }
        initialState.unitPath.add(h);
        pathStates.add(initialState);
        return pathStates;
    }

    public void calculateWithBFS(Context context, UnitGraph ug, List<PathState> pathStates, List<String> testSet) throws Exception {
        while(!pathStates.isEmpty()){
            PathState pathState = pathStates.remove(0);
            Unit unit = pathState.node;
            if(pathState.unitPath.size()>100||unit instanceof JReturnStmt||unit instanceof JReturnVoidStmt){
                if(pathCheck(pathState.unitPath)){
                    List<BoolExpr> pathConstraint = calPathConstraint(context, pathState.unitPath);
                    String testSetByPathConstraint = getTestSetByPathConstraint(context, pathConstraint);
                    if(testSetByPathConstraint!=null){
                        testSet.add(testSetByPathConstraint);
                        visitedNode.addAll(pathState.unitPath);
//                        System.out.println("111"+visitedNode);
                    }
                }
                continue;
            }

            Solver solver = context.mkSolver();
            List<BoolExpr> pathConstraint1 = calPathConstraint(context, pathState.unitPath);
            if(!pathConstraint1.isEmpty()) {
                for(BoolExpr expr : pathConstraint1 ){
                    solver.add(expr);
                }

                if(solver.check() == Status.UNSATISFIABLE){
                    continue;
                }
            }

            visitedNode.add(unit);
            if(visitedNode.size() == ug.size()){
                // 如果所有节点都访问过了，那么就是一条路径
                // 生成一个新的状态
                // 生成一个新的约束
                break;
            }


            if(unit instanceof JIfStmt){
                // 如果是if语句，那么就是两条路径
                // 生成两个新的状态
                PathState pathState1 = new PathState(pathState);
                PathState pathState2 = new PathState(pathState);
                // 生成两个新的约束
                // 有可能是一个循环，判断一下

                // BoolExpr ifstms = (BoolExpr) getZ3ExprFromSootExpr(context, pathState.params, ((JIfStmt) unit).getCondition());
                List<Unit> nextUnitList = ug.getSuccsOf(unit);
                Unit nextUnit1 = nextUnitList.get(0);
                Unit nextUnit2 = nextUnitList.get(1);
                pathState1.node = nextUnit1;
                pathState2.node = nextUnit2;
                pathState1.unitPath.add(nextUnit1);
                pathState2.unitPath.add(nextUnit2);
                // pathState1.PathConstraints.add(context.mkNot(ifstms));
                // pathState2.PathConstraints.add(ifstms);


                pathStates.add(pathState1);
                pathStates.add(pathState2);
            }else if(unit instanceof JAssignStmt){
                // 如果是赋值语句，那么就是一条路径
                // 生成一个新的状态
                PathState pathState1 = new PathState(pathState);
                // 生成一个新的约束
                String leftOp = ((JAssignStmt) unit).getLeftOp().toString();
                pathState1.params.put(leftOp, context.mkIntConst(leftOp));
                // Expr rightOp = getZ3ExprFromSootExpr(context, pathState1.params, ((JAssignStmt) unit).getRightOp());
                // if (rightOp == null) {
                //    System.out.println("Ignore the statement: " + unit.toString());
                //    continue;
                //}
                // pathState1.PathConstraints.add(context.mkEq(pathState1.params.get(leftOp), rightOp));
                pathState1.node = ug.getSuccsOf(unit).get(0);
                pathState1.unitPath.add(pathState1.node);
                pathStates.add(pathState1);

            }else if(unit instanceof JGotoStmt){
                // 如果是goto语句，那么就是一条路径
                // 生成一个新的状态
                PathState pathState1 = new PathState(pathState);
                // 判断是否是循环
                Unit nextUnit = ug.getSuccsOf(unit).get(0);

                pathState1.node = nextUnit;
                pathState1.unitPath.add(pathState1.node);

                pathStates.add(pathState1);
            }else {
                // 如果是其它语句，那么就是一条路径
                // 生成一个新的状态
                PathState pathState1 = new PathState(pathState);
                // 生成一个新的约束
                pathState1.node = ug.getSuccsOf(unit).get(0);
                pathState1.unitPath.add(pathState1.node);
                pathStates.add(pathState1);
            }
        }
    }

    private boolean pathCheck(List<Unit> unitPath) {
        boolean flag = false;
        for (Unit unit : unitPath) {
            if(visitedNode.contains(unit)){
                continue;
            }
            flag = true;
        }
        return flag;
    }

    public List<BoolExpr> calPathConstraint(Context context, List<Unit> path) {
        Map<String, Expr> paras = new HashMap<>();
        Map<String, Integer> latestAssign = new HashMap<>();
        Z3Parser parser = new Z3Parser(context, paras, latestAssign);
        List<BoolExpr> pathConstraint = new ArrayList<>();
        for (Local local : getParameters()) {
            latestAssign.put(local.getName(), 0);
            paras.put(local.getName() + "_0", parser.fromLocal(local, pathConstraint));
        }
        String expectedResult = "";

        for (int i = 0; i < path.size(); i++) {
            Unit stmt = path.get(i);

            if (stmt instanceof JAssignStmt) {
                // assignList.put(((JAssignStmt) stmt).getLeftOp().toString(), ((JAssignStmt) stmt).getRightOp().toString());
                // String leftOp = ((JAssignStmt) stmt).getLeftOp().toString();
                // paras.put(leftOp, context.mkIntConst(leftOp));
                Expr rightOp = parser.fromSootExpr(((JAssignStmt) stmt).getRightOp());
                if (rightOp == null) {
                    // leftOp 设置为对应类型任意值
                    parser.assignLocal((Local) ((JAssignStmt) stmt).getLeftOp(), pathConstraint);
                    System.err.println("Ignore the statement: " + stmt + " assign arbitrary value to " + ((JAssignStmt) stmt).getLeftOp().toString());
                    continue;
                }
                Value leftOp = ((JAssignStmt) stmt).getLeftOp();
                if (leftOp instanceof Local) {
                    Expr newLeft = parser.assignLocal((Local) leftOp, pathConstraint);
                    pathConstraint.add(context.mkEq(newLeft, rightOp));
                } else if (leftOp instanceof JArrayRef) {
                    Value base = ((JArrayRef) leftOp).getBase();
                    Value index = ((JArrayRef) leftOp).getIndex();
                    if (base instanceof Local) {
                        ArrayExpr rightStore = context.mkStore(parser.fromLocal((Local) base, pathConstraint), parser.fromSootExpr(index), rightOp);
                        ArrayExpr newLeft = (ArrayExpr) parser.assignLocal((Local) base, pathConstraint);
                        pathConstraint.add(context.mkEq(newLeft, rightStore));
                    } else {
                        System.out.println("Ignore the statement: " + stmt.toString());
                        continue;
                    }
                } else {
                    System.out.println("Ignore the statement: " + stmt.toString());
                    continue;
                }
                continue;
            }
            if (stmt instanceof JIfStmt) {
//                System.out.println("The ifstmt is: " + stmt);
                BoolExpr ifstms = (BoolExpr) parser.fromSootExpr(((JIfStmt) stmt).getCondition());
                Unit nextUnit = i+1 == path.size() ? null : path.get(i + 1);
//                System.out.println("The nextUnit is: " + nextUnit);

                //如果ifstmt的后继语句不是ifstmt中goto语句，说明ifstmt中的条件为假
                if (!((JIfStmt) stmt).getTarget().equals(nextUnit))
                    ifstms = context.mkNot(ifstms);
                pathConstraint.add(ifstms);
                continue;
            }
            if (stmt instanceof JReturnStmt) {
                expectedResult = stmt.toString().replace("return", "").trim();
            }
        }
//        System.out.println("The step conditions with JimpleVars are: " + pathConstraint);
        return pathConstraint;
    }


}
